import unittest, operator
import numpy as np
from .. import cv, data, perf
from ..nodes import PriorClassifier
from .. import DataSet

class TestFolds(unittest.TestCase):
  def setUp(self):
    self.d = data.gaussian_dataset([30, 20, 10])

  def test_strat_splits(self):
    self.assertRaises(AssertionError, cv.strat_splits, self.d.Y, 31)

  def test_seq_splits(self):
    self.assertRaises(AssertionError, cv.seq_splits, 30, 31)

  def test_subsets(self):
    fis = (np.arange(self.d.ninstances) % 3) - 1 # test negative fis
    ds = cv.subsets(self.d, fis)
    for (i, fi) in enumerate(np.unique(fis)):
      assert ds[i] == self.d[fis==fi]

  def test_folds(self):
    fis = (np.arange(self.d.ninstances) % 3)
    subs = cv.subsets(self.d, fis)
    folds = cv.folds(subs)
    for fi, (tr, te) in enumerate(folds):
      assert te == subs[fi]
      assert tr == reduce(operator.add, [subs[i] for i in range(3) if i != fi]) 
      

class TestCrossValidation(unittest.TestCase):
  def setUp(self):
    self.d = data.gaussian_dataset([30, 20, 10])

  def test_crossvalidation_label_prot(self):
    '''Test protection against reading labels in cross-validation'''
    class CheckNode:
      def train_apply(self, tr, te): 
        assert(np.all(np.isnan(te.Y)))
        return te

    subs = cv.subsets(self.d, cv.strat_splits(self.d.Y, 4))
    preds = cv.cross_validate(cv.folds(subs), CheckNode())

  def test_crossvalidation_mem_prot(self):
    '''Test protection against remembering over repetitions'''
    class MemNode:
      def __init__(self):
        self.mem = {}

      def train_apply(self, tr, te):
        # train
        self.mem.update(zip(tr.I[0], tr.Y.T.tolist()))
        # apply
        X = np.asarray([self.mem.get(i, [0, 0, 0]) for i in te.I[0]]).T
        return DataSet(X=X, default=te)

    subs = cv.subsets(self.d, cv.strat_splits(self.d.Y, 4))
    a = perf.mean_std(perf.accuracy, 
      cv.cross_validate(cv.folds(subs), MemNode()))
    self.assertAlmostEqual(a[0], .5, 1)

  def test_rep_cv(self):
    tds = cv.rep_cv(self.d, PriorClassifier(), k=10, reps=5)
    self.assertEqual(len(tds), 50)
    self.assertAlmostEqual(perf.mean_std(perf.accuracy, tds)[0], .5)
    rep_ids = np.array([td.I.flatten() for td in tds])
    rep_ids = rep_ids.reshape(-1, self.d.ninstances)
    self.failIf((np.var(rep_ids, axis=0) == 0).any(),
      'The folds are all the same!')
